//
//  CTAesUtil.m
//  CTOpenSSL
//
//  Created by Andney Jack on 2021/2/22.
//

#import "CTAesUtil.h"
#include <openssl/aes.h>

@implementation CTAesUtil

+ (NSString *)aesEncode:(id)content key:(NSString *)key initVector:(NSString *)initVector
{
    NSData *data = nil;
    if (content && [content isKindOfClass:[NSString class]]) {
        data = [(NSString *)content dataUsingEncoding:NSUTF8StringEncoding];
    } else if (content && [content isKindOfClass:[NSData class]]) {
        data = content;
    }

    if (!data || data.length < 1) {
        [self m_logMsg:@"source data is nil or empty."];
        return nil;
    }

    if (key.length != AES_BLOCK_SIZE) {
        [self m_logMsg:[NSString stringWithFormat:@"AES加密的key长度不是%d位", AES_BLOCK_SIZE]];
        return nil;
    }
    if (initVector.length != AES_BLOCK_SIZE) {
        [self m_logMsg:[NSString stringWithFormat:@"AES加密的初始化向量iv长度不是%d位", AES_BLOCK_SIZE]];
        return nil;
    }

    size_t srcDataLength = data.length; // 源数据长度
    size_t destDataLength = 0; // 加密后的数据长度
    int fillNumber = 0; // 填充数字
    if ((srcDataLength % AES_BLOCK_SIZE) == 0) {
        destDataLength = srcDataLength + AES_BLOCK_SIZE;
        fillNumber = AES_BLOCK_SIZE;
    } else {
        destDataLength = ((srcDataLength / AES_BLOCK_SIZE) + 1) * AES_BLOCK_SIZE;
        fillNumber = (int)(destDataLength - srcDataLength);
    }

    unsigned char *inputData = (unsigned char *)malloc(destDataLength + 1);
    if (inputData == NULL) {
        [self m_logMsg:@"Unable to allocate memory for inputData"];
        return nil;
    }

    // 填充，PKCS5填充的规则：如果加密数据余16大于0 将余数填充到加密数据后面凑成待加密数据长度正好是16的整数倍，如果加密数据余16正好等于0，那么在加密数据后面填充16个16
    memset(inputData, fillNumber, destDataLength);
    inputData[destDataLength] = 0x00;
    memcpy(inputData, data.bytes, srcDataLength); // 拷贝资源

    // 开辟加密后的数据空间
    unsigned char *encryptBuffer = (unsigned char *)malloc(destDataLength + 1);
    if (encryptBuffer == NULL) {
        [self m_logMsg:@"Unable to allocate memory for encryptBuffer"];
        return nil;
    }
    memset(encryptBuffer, 0, destDataLength + 1);

    // 设置key，iv
    unsigned char keyBuffer[AES_BLOCK_SIZE + 1];
    unsigned char ivBuffer[AES_BLOCK_SIZE + 1];

    memset(keyBuffer, 0x00, sizeof(keyBuffer));
    memcpy(keyBuffer, [key cStringUsingEncoding:NSUTF8StringEncoding], AES_BLOCK_SIZE);
    memset(ivBuffer, 0x00, sizeof(ivBuffer));
    memcpy(ivBuffer, [initVector cStringUsingEncoding:NSUTF8StringEncoding], AES_BLOCK_SIZE);

    // 设置加密密钥
    AES_KEY aesKey;
    memset(&aesKey, 0x00, sizeof(AES_KEY));
    if (AES_set_encrypt_key(keyBuffer, 128, &aesKey) < 0) {
        [self m_logMsg:@"Unable to set encryption key in AES..."];
        return nil;
    }

    // 加密
    AES_cbc_encrypt((unsigned char *)inputData, (unsigned char *)encryptBuffer,
                    destDataLength, &aesKey, ivBuffer, AES_ENCRYPT);

    // 将加密后的数据转成NSData
    NSData *result = [NSData dataWithBytes:encryptBuffer length:destDataLength];

    // 释放内存
    if (inputData != NULL) {
        free(inputData);
        inputData = NULL;
    }
    if (encryptBuffer != NULL) {
        free(encryptBuffer);
        encryptBuffer = NULL;
    }

    NSString *base64 = [result base64EncodedStringWithOptions:0];
    return base64 && base64.length > 0 ? base64 : nil;
}

+ (NSData *)aesDecodeToData:(id)content key:(NSString *)key initVector:(NSString *)initVector
{
    NSData *base64Data = nil;
    if (content && [content isKindOfClass:[NSString class]]) {
        base64Data = [[NSData alloc] initWithBase64EncodedString:content options:NSDataBase64DecodingIgnoreUnknownCharacters];
    } else if (content && [content isKindOfClass:[NSData class]]) {
        base64Data = content;
    }

    if (!base64Data || base64Data.length < 1) {
        [self m_logMsg:@"source data is nil or empty."];
        return nil;
    }

    if (key.length != AES_BLOCK_SIZE) {
        [self m_logMsg:[NSString stringWithFormat:@"AES解密的key长度不是%d位", AES_BLOCK_SIZE]];
        return nil;
    }
    if (initVector.length != AES_BLOCK_SIZE) {
        [self m_logMsg:[NSString stringWithFormat:@"AES解密的初始化向量iv长度不是%d位", AES_BLOCK_SIZE]];
        return nil;
    }

    size_t srcDataLength = base64Data.length; // 源数据长度
    size_t destDataLength = 0; // 解密后的数据长度

    if ((srcDataLength % AES_BLOCK_SIZE) == 0) {
        destDataLength = srcDataLength;
    } else {
        destDataLength = ((srcDataLength / AES_BLOCK_SIZE) + 1) * AES_BLOCK_SIZE;
    }

    // 设置key，iv
    unsigned char keyBuffer[AES_BLOCK_SIZE + 1];
    unsigned char ivBuffer[AES_BLOCK_SIZE + 1];

    memset(keyBuffer, 0x00, sizeof(keyBuffer));
    memcpy(keyBuffer, [key cStringUsingEncoding:NSUTF8StringEncoding], AES_BLOCK_SIZE);
    memset(ivBuffer, 0x00, sizeof(ivBuffer));
    memcpy(ivBuffer, [initVector cStringUsingEncoding:NSUTF8StringEncoding], AES_BLOCK_SIZE);

    // 设置解密密钥
    AES_KEY aesKey;
    memset(&aesKey, 0x00, sizeof(AES_KEY));
    if (AES_set_decrypt_key(keyBuffer, 128, &aesKey) < 0) {
        NSLog(@"Unable to set decryption key in AES...");
        return nil;
    }

    unsigned char *decryptBuffer = (unsigned char *)calloc(destDataLength + AES_BLOCK_SIZE, sizeof(unsigned char));
    if (decryptBuffer == NULL) {
        [self m_logMsg:@"Unable to allocate memory for decryptBuffer"];
        return nil;
    }

    // 解密
    AES_cbc_encrypt((unsigned char *)base64Data.bytes, decryptBuffer,
                    destDataLength, &aesKey, ivBuffer, AES_DECRYPT);

    unsigned char fillNumber = decryptBuffer[destDataLength - 1];
    NSInteger targetLength = destDataLength - fillNumber;

    if (targetLength < 1) {
        // 释放内存
        if (decryptBuffer != NULL) {
            free(decryptBuffer);
            decryptBuffer = NULL;
        }
        return nil;
    }

    NSData *result = [NSData dataWithBytes:decryptBuffer length:targetLength];
    // 释放内存
    if (decryptBuffer != NULL) {
        free(decryptBuffer);
        decryptBuffer = NULL;
    }

    return result;
}

+ (NSString *)aesDecodeToString:(id)content key:(NSString *)key initVector:(NSString *)initVector {
    NSData *data = [self aesDecodeToData:content key:key initVector:initVector];
    if (data && data.length > 0) {
        return [[NSString alloc] initWithData:data encoding:NSUTF8StringEncoding];
    }
    return nil;
}

+ (void)m_logMsg:(NSString *)msg
{
#ifdef DEBUG
    NSLog(@"[CTAesUtil]: %@", msg);
#endif
}

@end
